

from langchain.agents import tool,AgentExecutor
from langchain.agents.format_scratchpad.openai_tools import format_to_openai_tool_messages
from langchain.agents.output_parsers.openai_tools import OpenAIToolsAgentOutputParser
from langchain.chat_models import ChatOpenAI
from langchain.prompts import MessagesPlaceholder
from langchain_community.tools.convert_to_openai import format_tool_to_openai_tool
from langchain_core.messages import AIMessage, HumanMessage
from langchain_core.prompts import ChatPromptTemplate
import time
import jwt


ZHIPU_KEY = "ab5070471f43c2267d9de50d91e3ba0b.IlnmhMI3RnJlY1Zc"

def generate_token(apikey: str, exp_seconds=1000000):
    try:
        id, secret = apikey.split(".")
    except Exception as e:
        raise Exception("invalid apikey", e)

    payload = {
        "api_key": id,
        "exp": int(round(time.time() * 1000)) + exp_seconds * 1000,
        "timestamp": int(round(time.time() * 1000)),
    }

    return jwt.encode(
        payload,
        secret,
        algorithm="HS256",
        headers={"alg": "HS256", "sign_type": "SIGN"},
    )


# 定义大模型
def get_glm(temprature=0.95):
    llm = ChatOpenAI(
        model_name="glm-4",
        openai_api_base="https://open.bigmodel.cn/api/paas/v4",
        openai_api_key=generate_token(ZHIPU_KEY),
        streaming=False,
        temperature=temprature
    )
    return llm


# 定义Tools
# TODO:定义一个工具函数，用于获取单词的长度
@tool
def get_word_length(word: str) -> int:
    return len(word)


tools = [get_word_length]

# 定义prompt
prompt = ChatPromptTemplate.from_messages(
    [
        (
            "system",
            "你是个非常强大的助手，但是不善于做数学计算",
        ),
        ("user", "{input}"),
        MessagesPlaceholder(variable_name="agent_scratchpad"),
    ]
)

# 绑定Tool到LLM
# 把zhipu的模型转换到OpenAI的格式。通过Langchain
llm_with_tools = get_glm(0.01).bind(tools=[format_tool_to_openai_tool(tool) for tool in tools])

# 创建Agent
agent = (
        {
            "input": lambda x: x["input"],
            "agent_scratchpad": lambda x: format_to_openai_tool_messages(
                x["intermediate_steps"]
            ),
        }
        | prompt
        | llm_with_tools
        | OpenAIToolsAgentOutputParser()
)

agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)

# 增加Memory
MEMORY_KEY = "chat_history"
prompt = ChatPromptTemplate.from_messages(
    [
        (
            "system",
            "你是个非常强大的助手，但是不善于做数学计算",
        ),
        MessagesPlaceholder(variable_name=MEMORY_KEY),
        ("user", "{input}"),
        MessagesPlaceholder(variable_name="agent_scratchpad"),
    ]
)

# 用List来跟踪会话历史，并用AIMessage， HumanMessage来简化Message的构建。
chat_history = []
agent = (
        {
            "input": lambda x: x["input"],
            "agent_scratchpad": lambda x: format_to_openai_tool_messages(
                x["intermediate_steps"]
            ),
            "chat_history": lambda x: x["chat_history"],
        }
        | prompt
        | llm_with_tools
        | OpenAIToolsAgentOutputParser()
)
agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)




if __name__ == "__main__":
    input1 = "下面这个单词有多少字母：eudcaaaabbbbcccceeeeeddddd"
    result = agent_executor.invoke({"input": input1, "chat_history": chat_history})
    chat_history.extend(
        [
            HumanMessage(content=input1),
            AIMessage(content=result["output"]),
        ]
    )